"""
@Filename       : dblp_preprocess.py
@Create Time    : 2021/9/23 9:40
@Author         : Rylynn
@Description    : 

"""
import random


def parse_paper(text_block, paper_dict):
    paper_meta = text_block.split('\n')
    references = []
    paper_entity = {}
    for line in paper_meta:
        line = line.strip().replace('\n', '')
        if line[:2] == '#*':
            paper_name = line.replace('#*', '')
            paper_entity['title'] = paper_name
        elif line[:2] == '#@':
            authors = list(map(lambda author: author.strip(), line.replace('#@', '').split(',')))
            paper_entity['authors'] = authors[0]
        elif line[:2] == '#i':
            idx = line.replace('#index', '')
            paper_entity['idx'] = idx
        elif line[:2] == '#t':
            year = line.replace('#t', '')
            paper_entity['year'] = eval(year)
        elif line[:2] == '#%':
            references.append(line.replace('#%', ''))

    references = set(references)
    paper_entity['ref'] = references

    if 'year' in paper_entity.keys() and 'title' in paper_entity.keys() and 'authors' in paper_entity.keys()\
            and paper_entity['year'] <= 2016:
        if paper_dict.get(paper_entity['idx']):
            exist_paper = paper_dict.get(paper_entity['idx'])
            if len(paper_entity['ref']) > len(exist_paper['ref']):
                paper_dict[paper_entity['idx']] = paper_entity
        else:
            paper_dict[paper_entity['idx']] = paper_entity
    else:
        return

def parse_all_paper(file_path):
    print("Extracting paper entity from file...")
    paper_dict = {}
    with open(file_path, encoding='utf8') as fp:
        text_block = ''
        block_line_count = 0
        for line in fp.readlines():
            if line == '\n':
                if block_line_count <= 0:
                    ...
                else:
                    parse_paper(text_block, paper_dict)
                text_block = ''
                block_line_count = 0
                continue
            text_block = text_block + line
            block_line_count += 1
    print("Extraction finish, totally {} paper.".format(len(paper_dict.keys())))
    return paper_dict


def build_cascade(paper_dict):
    author_dict = {}
    max_author_id = 1
    print("Building cascade from paper entity...")
    cascade_dict = {}
    for idx, paper_entity in paper_dict.items():
        references = paper_entity['ref']
        for ref in references:
            if ref not in paper_dict.keys():
                continue
            if cascade_dict.get(ref):
                cascade_dict[ref].append((idx, paper_entity['year']))
            else:
                cascade_dict[ref] = [(idx, paper_entity['year'])]

    for paper_id, citations in cascade_dict.items():
        author_list = []
        author_set = set()
        citations = sorted(citations, key=lambda c: c[1], reverse=False)
        citations = list(map(lambda c: c[0], citations))
        for c in citations:
            author = paper_dict[c]['authors']
            if author in author_set:
                continue
            else:
                author_set.add(author)
                if not author_dict.get(author):
                    author_dict[author] = max_author_id
                    max_author_id += 1
            author_list.append(author_dict[author])

        cascade_dict[paper_id] = author_list
    user_dict = {}
    for c in cascade_dict.values():
        for user in c:
            if user in user_dict.keys():
                user_dict[user] += 1
            else:
                user_dict[user] = 1

    active_user = set([k for k, v in user_dict.items() if v > 70])

    new_cascade_dict = {}
    for paper_id, cascade in cascade_dict.items():
        new_cascade = []
        for user in cascade:
            if user in active_user:
                new_cascade.append(user)
        if 10 < len(new_cascade) < 500:
            new_cascade_dict[paper_id] = new_cascade

    new_user_set = set()
    for c in new_cascade_dict.values():
        for user in c:
            new_user_set.add(user)

    print("Building finish, totally {} cascade.".format(len(new_cascade_dict.keys())))
    print("User number: ", len(new_user_set))

    with open('../../../data/dblp_new/content.txt', 'w', encoding='utf8') as f:
        for paper_id, _ in new_cascade_dict.items():
            f.write("{}\t{}\n".format(paper_id, paper_dict[paper_id]['title']))

    # final_cascade = list(new_cascade_dict.items())
    # random.shuffle(final_cascade)
    #
    # train_cascade = final_cascade[:int(len(final_cascade) * 0.8)]
    # test_cascade = final_cascade[int(len(final_cascade) * 0.8): int(len(final_cascade) * 0.9)]
    # valid_cascade = final_cascade[int(len(final_cascade) * 0.9):]
    #
    # def write_cascade(cascade, filename):
    #     with open(filename, 'w') as f:
    #         for k, c in cascade:
    #             f.write('{} '.format(k))
    #             for u in c:
    #                 f.write('{},{} '.format(u, 0))
    #             f.write('\n')
    #
    # write_cascade(train_cascade, '../../../data/dblp_new/cascade.txt')
    # write_cascade(test_cascade, '../../../data/dblp_new/cascadetest.txt')
    # write_cascade(valid_cascade, '../../../data/dblp_new/cascadevalid.txt')

if __name__ == '__main__':
    paper_dict = parse_all_paper('../../../data/raw/dblp.txt')
    build_cascade(paper_dict)
